import torch
import torch.nn as nn
class CrowdCounter(nn.Module):
    def __init__(self,model_name):
        super(CrowdCounter, self).__init__()
        if model_name == 'VGG':
            from .net.VGG import VGG as net
        elif model_name == 'MCNN':
            from .net.MCNN import mcnn as net

        self.CCN = net()

    @property
    def loss(self):
        return self.loss_mse

    def forward(self, img):
        density_map = self.CCN(img)
        return density_map

    def build_loss(self, density_map, gt_data):
        loss_mse = self.loss_mse_fn(density_map, gt_data)
        return loss_mse

    def test_forward(self, img):
        density_map = self.CCN(img)
        return density_map
